import Foundation
import Accelerate

// MARK: - MFCC特征提取器

final class MFCCExtractor {
    
    // MARK: - 配置参数（匹配训练参数）
    
    private let sampleRate: Int = 25        // Hz
    private let numCepstral: Int = 13       // MFCC系数数量
    private let numMelFilters: Int = 26     // 梅尔滤波器数量
    private let fftSize: Int = 64           // FFT大小（匹配窗口大小）
    private let windowSize: Int = 64        // 窗口大小
    
    // 预计算的梅尔滤波器组
    private var melFilterbank: [[Float]]
    
    // 预计算的DCT矩阵
    private var dctMatrix: [[Float]]
    
    // MARK: - 初始化
    
    init() {
        // 预计算梅尔滤波器组
        self.melFilterbank = MFCCExtractor.computeMelFilterbank(
            numFilters: numMelFilters,
            fftSize: fftSize,
            sampleRate: Float(sampleRate),
            lowFreq: 0.5,   // 匹配训练时的lowcut
            highFreq: 6.0   // 匹配训练时的highcut
        )
        
        // 预计算DCT矩阵
        self.dctMatrix = MFCCExtractor.computeDCTMatrix(
            numCepstral: numCepstral,
            numFilters: numMelFilters
        )
    }
    
    // MARK: - 主要方法
    
    /// 从三轴加速度提取117维MFCC特征
    func extractFeatures(
        ax: [Double],
        ay: [Double],
        az: [Double]
    ) -> [Float] {
        guard ax.count == windowSize,
              ay.count == windowSize,
              az.count == windowSize else {
            print("❌ [MFCC] 窗口大小不匹配: \(ax.count) vs \(windowSize)")
            return Array(repeating: 0, count: 117)
        }
        
        var allFeatures: [Float] = []
        
        // 对每个轴提取MFCC特征
        for buffer in [ax, ay, az] {
            // 1. 基础MFCC (13维)
            let mfcc = computeMFCC(buffer.map { Float($0) })
            
            // 2. 一阶差分 (13维)
            let delta = computeDelta(mfcc)
            
            // 3. 二阶差分 (13维)
            let deltaDelta = computeDelta(delta)
            
            // 合并: 13 + 13 + 13 = 39维/轴
            allFeatures.append(contentsOf: mfcc)
            allFeatures.append(contentsOf: delta)
            allFeatures.append(contentsOf: deltaDelta)
        }
        
        // 总计: 39 × 3 = 117维
        return allFeatures
    }
    
    // MARK: - MFCC计算
    
    /// 计算MFCC系数
    private func computeMFCC(_ signal: [Float]) -> [Float] {
        // 1. 预加重（可选，增强高频）
        let preemphasized = applyPreemphasis(signal, coefficient: 0.97)
        
        // 2. 加窗（汉明窗）
        let windowed = applyHammingWindow(preemphasized)
        
        // 3. FFT
        let powerSpectrum = computePowerSpectrum(windowed)
        
        // 4. 梅尔滤波器组
        let melEnergies = applyMelFilterbank(powerSpectrum)
        
        // 5. 对数
        let logMelEnergies = melEnergies.map { log(max($0, 1e-10)) }
        
        // 6. DCT
        let mfcc = applyDCT(logMelEnergies)
        
        return mfcc
    }
    
    /// 预加重
    private func applyPreemphasis(_ signal: [Float], coefficient: Float) -> [Float] {
        guard signal.count > 1 else { return signal }
        
        var output = [Float](repeating: 0, count: signal.count)
        output[0] = signal[0]
        
        for i in 1..<signal.count {
            output[i] = signal[i] - coefficient * signal[i-1]
        }
        
        return output
    }
    
    /// 汉明窗
    private func applyHammingWindow(_ signal: [Float]) -> [Float] {
        let N = signal.count
        var output = [Float](repeating: 0, count: N)
        
        for i in 0..<N {
            let window = 0.54 - 0.46 * cos(2.0 * Float.pi * Float(i) / Float(N - 1))
            output[i] = signal[i] * window
        }
        
        return output
    }
    
    /// 计算功率谱
    private func computePowerSpectrum(_ signal: [Float]) -> [Float] {
        var input = signal
        
        // 补零到fftSize
        while input.count < fftSize {
            input.append(0)
        }
        
        // FFT
        var realPart = [Float](repeating: 0, count: fftSize / 2)
        var imagPart = [Float](repeating: 0, count: fftSize / 2)
        
        input.withUnsafeMutableBufferPointer { inputPtr in
            var splitComplex = DSPSplitComplex(
                realp: &realPart,
                imagp: &imagPart
            )
            
            inputPtr.baseAddress!.withMemoryRebound(to: DSPComplex.self, capacity: fftSize / 2) { complexPtr in
                vDSP_ctoz(complexPtr, 2, &splitComplex, 1, vDSP_Length(fftSize / 2))
            }
            
            // 执行FFT
            let log2n = vDSP_Length(log2(Float(fftSize)))
            guard let fftSetup = vDSP_create_fftsetup(log2n, FFTRadix(kFFTRadix2)) else {
                return
            }
            
            vDSP_fft_zrip(fftSetup, &splitComplex, 1, log2n, FFTDirection(FFT_FORWARD))
            
            vDSP_destroy_fftsetup(fftSetup)
        }
        
        // 计算功率谱: |X|^2
        var powerSpectrum = [Float](repeating: 0, count: fftSize / 2)
        var splitComplex = DSPSplitComplex(realp: &realPart, imagp: &imagPart)
        vDSP_zvmags(&splitComplex, 1, &powerSpectrum, 1, vDSP_Length(fftSize / 2))
        
        return powerSpectrum
    }
    
    /// 应用梅尔滤波器组
    private func applyMelFilterbank(_ powerSpectrum: [Float]) -> [Float] {
        var melEnergies = [Float](repeating: 0, count: numMelFilters)
        
        for (i, filter) in melFilterbank.enumerated() {
            var energy: Float = 0
            for (j, weight) in filter.enumerated() {
                if j < powerSpectrum.count {
                    energy += weight * powerSpectrum[j]
                }
            }
            melEnergies[i] = energy
        }
        
        return melEnergies
    }
    
    /// 应用DCT变换
    private func applyDCT(_ melEnergies: [Float]) -> [Float] {
        var mfcc = [Float](repeating: 0, count: numCepstral)
        
        for i in 0..<numCepstral {
            var sum: Float = 0
            for (j, energy) in melEnergies.enumerated() {
                sum += dctMatrix[i][j] * energy
            }
            mfcc[i] = sum
        }
        
        return mfcc
    }
    
    /// 计算差分
    private func computeDelta(_ coeffs: [Float], N: Int = 2) -> [Float] {
        let numFrames = 1  // 我们只有一帧
        guard numFrames == 1 else {
            // 如果有多帧，使用标准差分公式
            return coeffs  // 简化处理
        }
        
        // 单帧情况，返回零差分
        return [Float](repeating: 0, count: coeffs.count)
    }
    
    // MARK: - 静态辅助方法
    
    /// 计算梅尔滤波器组
    static func computeMelFilterbank(
        numFilters: Int,
        fftSize: Int,
        sampleRate: Float,
        lowFreq: Float,
        highFreq: Float
    ) -> [[Float]] {
        // 频率到梅尔
        func hzToMel(_ hz: Float) -> Float {
            return 2595.0 * log10(1.0 + hz / 700.0)
        }
        
        // 梅尔到频率
        func melToHz(_ mel: Float) -> Float {
            return 700.0 * (pow(10.0, mel / 2595.0) - 1.0)
        }
        
        // 梅尔刻度上的均匀间隔点
        let lowMel = hzToMel(lowFreq)
        let highMel = hzToMel(highFreq)
        let melPoints = (0...numFilters + 1).map { i in
            lowMel + Float(i) * (highMel - lowMel) / Float(numFilters + 1)
        }
        
        // 转回Hz并映射到FFT bins
        let binPoints = melPoints.map { mel in
            Int(floor(Float(fftSize + 1) * melToHz(mel) / sampleRate))
        }
        
        // 构建滤波器
        var filterbank = [[Float]](repeating: [Float](repeating: 0, count: fftSize / 2),
                                   count: numFilters)
        
        for i in 0..<numFilters {
            let left = binPoints[i]
            let center = binPoints[i + 1]
            let right = binPoints[i + 2]
            
            // 上升斜坡
            for j in left..<center {
                if j < fftSize / 2 {
                    let denominator = center - left
                    if denominator > 0 {
                        filterbank[i][j] = Float(j - left) / Float(denominator)
                    }
                }
            }
            
            // 下降斜坡
            for j in center..<right {
                if j < fftSize / 2 {
                    let denominator = right - center
                    if denominator > 0 {
                        filterbank[i][j] = Float(right - j) / Float(denominator)
                    }
                }
            }
        }
        
        return filterbank
    }
    
    /// 计算DCT矩阵
    static func computeDCTMatrix(numCepstral: Int, numFilters: Int) -> [[Float]] {
        var matrix = [[Float]](repeating: [Float](repeating: 0, count: numFilters),
                              count: numCepstral)
        
        let norm = sqrt(2.0 / Float(numFilters))
        
        for i in 0..<numCepstral {
            for j in 0..<numFilters {
                matrix[i][j] = norm * cos(Float.pi * Float(i) * (Float(j) + 0.5) / Float(numFilters))
            }
        }
        
        return matrix
    }
}
